#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import json
import glob
import cv2
from tqdm import tqdm
from PIL import Image
import argparse
import numpy as np
from datetime import datetime

def yolo2coco(yolo_dataset_path, output_dir):
    """
    将YOLO格式的数据集转换为COCO格式，并为每个子集生成独立的JSON文件
    
    Args:
        yolo_dataset_path: YOLO数据集的根目录
        output_dir: 输出的COCO格式JSON文件的目录
    """
    # 获取类别信息
    classes_file = os.path.join(yolo_dataset_path, "classes.txt")
    
    if not os.path.exists(classes_file):
        raise FileNotFoundError(f"未找到classes.txt文件")
    
    # 读取类别信息
    with open(classes_file, 'r', encoding='utf-8') as f:
        classes = [line.strip() for line in f.readlines() if line.strip()]
    
    # 创建COCO格式的数据结构
    coco_data = {
        "info": {
            "description": "PCB_DATASET",
            "url": "",
            "version": "1.0",
            "year": datetime.now().year,
            "contributor": "",
            "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        },
        "licenses": [
            {
                "id": 1,
                "name": "Unknown",
                "url": ""
            }
        ],
        "categories": [],
        "images": [],
        "annotations": []
    }
    
    # 添加类别信息 - 从1开始
    for i, cls_name in enumerate(classes):
        coco_data["categories"].append({
            "id": i + 1,  # 类别ID从1开始
            "name": cls_name,
            "supercategory": "PCB"
        })
    
    # 创建类别名称到ID的映射
    class_name_to_id = {cls_name: i + 1 for i, cls_name in enumerate(classes)}
    
    # 处理每个子集（train, val, test）
    subsets = ["train", "val", "test"]
    for subset in subsets:
        subset_coco_data = {
            "info": coco_data["info"].copy(),
            "licenses": coco_data["licenses"].copy(),
            "categories": coco_data["categories"].copy(),
            "images": [],
            "annotations": []
        }
        
        # 获取子集的图像和标签文件路径
        image_dir = os.path.join(yolo_dataset_path, "images", subset)
        label_dir = os.path.join(yolo_dataset_path, "labels", subset)
        
        if not os.path.exists(image_dir):
            print(f"未找到子集 {subset} 的图像目录，跳过...")
            continue
        
        # 获取所有图像文件
        image_files = glob.glob(os.path.join(image_dir, "*.jpg")) + \
                     glob.glob(os.path.join(image_dir, "*.jpeg")) + \
                     glob.glob(os.path.join(image_dir, "*.png"))
        
        if not image_files:
            print(f"未找到子集 {subset} 的任何图像文件，跳过...")
            continue
        
        print(f"处理子集 {subset}，找到 {len(image_files)} 个图像文件")
        
        # 处理每个图像和对应的标签
        ann_id = 0
        for img_id, img_path in enumerate(tqdm(image_files, desc=f"处理 {subset} 图像")):
            # 获取图像信息
            img_filename = os.path.basename(img_path)
            img_name_without_ext = os.path.splitext(img_filename)[0]
            
            # 获取图像尺寸
            try:
                img = Image.open(img_path)
                img_width, img_height = img.size
            except Exception as e:
                print(f"无法打开图像 {img_path}: {e}")
                continue
            
            # 添加图像信息到COCO数据
            subset_coco_data["images"].append({
                "id": img_id + 1,  # 图像ID从1开始
                "file_name": img_filename,
                "width": img_width,
                "height": img_height,
                "license": 1,
                "date_captured": ""
            })
            
            # 确定标签文件路径
            label_path = os.path.join(label_dir, f"{img_name_without_ext}.txt")
            
            if not os.path.exists(label_path):
                print(f"未找到图像 {img_filename} 对应的标签文件")
                continue
            
            # 读取标签信息
            try:
                with open(label_path, 'r') as f:
                    label_lines = [line.strip() for line in f.readlines() if line.strip()]
                
                # 处理每个标签
                for line in label_lines:
                    parts = line.split()
                    if len(parts) != 5:
                        print(f"标签格式错误: {line}")
                        continue
                    
                    # 解析YOLO格式的标签
                    cls_id = int(parts[0])
                    if cls_id >= len(classes):
                        print(f"警告: 标签中的类别ID {cls_id} 超出了类别列表范围 (0-{len(classes)-1})")
                        continue
                    
                    # 获取类别名称和新的类别ID
                    cls_name = classes[cls_id]
                    new_cls_id = class_name_to_id[cls_name]  # 使用从1开始的类别ID
                    
                    x_center = float(parts[1])
                    y_center = float(parts[2])
                    width = float(parts[3])
                    height = float(parts[4])
                    
                    # 转换为COCO格式（左上角坐标和宽高）
                    x = (x_center - width / 2) * img_width
                    y = (y_center - height / 2) * img_height
                    w = width * img_width
                    h = height * img_height
                    
                    # 添加标注信息到COCO数据
                    subset_coco_data["annotations"].append({
                        "id": ann_id + 1,  # 标注ID从1开始
                        "image_id": img_id + 1,  # 使用图像ID
                        "category_id": new_cls_id,  # 使用从1开始的类别ID
                        "bbox": [x, y, w, h],
                        "area": w * h,
                        "segmentation": [],
                        "iscrowd": 0
                    })
                    ann_id += 1
            except Exception as e:
                print(f"处理标签文件 {label_path} 时出错: {e}")
        
        # 保存COCO格式的数据到JSON文件
        output_json_path = os.path.join(output_dir, f"{subset}.json")
        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(subset_coco_data, f, ensure_ascii=False, indent=2)
        
        print(f"子集 {subset} 转换完成，共处理 {len(subset_coco_data['images'])} 张图像，{len(subset_coco_data['annotations'])} 个标注")
        print(f"COCO格式数据已保存到: {output_json_path}")

def main():
    parser = argparse.ArgumentParser(description="将YOLO格式的数据集转换为COCO格式")
    parser.add_argument("--dataset_path", type=str, default=".", help="YOLO数据集的根目录")
    parser.add_argument("--output_dir", type=str, default=".", help="输出的COCO格式JSON文件的目录")
    args = parser.parse_args()
    
    yolo2coco(args.dataset_path, args.output_dir)

if __name__ == "__main__":
    main()
